import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.preprocessing import LabelEncoder
from scipy import stats
import os

plt.rcParams['font.sans-serif'] = ['SimHei']  # 显示中文
plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号
sns.set(style="whitegrid", font="SimHei", font_scale=1.1)

# 输出目录
output_dir = "../fig"
os.makedirs(output_dir, exist_ok=True)

# ========================
# 读取数据
# ========================
data = pd.read_csv('../data/train.csv')

target = 'Attrition'
categorical_features = [
    'BusinessTravel', 'Department', 'EducationField', 'Gender',
    'JobRole', 'MaritalStatus', 'Over18', 'OverTime'
]

results = []

for col in categorical_features:
    unique_vals = data[col].nunique()

    if unique_vals == 1:
        results.append((col, '单一取值', np.nan, np.nan))
        continue

    y = data[target]

    # 二分类 → 点双列相关
    if unique_vals == 2:
        le = LabelEncoder()
        x = le.fit_transform(data[col])
        r, p = stats.pointbiserialr(x, y)
        results.append((col, '点双列相关', r, p))

    # 多分类 → ANOVA
    else:
        groups = [group[target].values for name, group in data.groupby(col)]
        f_stat, p_value = stats.f_oneway(*groups)
        results.append((col, 'ANOVA', f_stat, p_value))

# ========================
# 结果整理
# ========================
result_df = pd.DataFrame(results, columns=['特征', '检验方法', '统计值/相关系数', 'p值'])
result_df['显著相关'] = result_df['p值'] < 0.05

# 排序（按统计值或相关系数强度）
result_df['绝对值'] = result_df['统计值/相关系数'].abs()
result_df = result_df.sort_values(by='绝对值', ascending=False)

# 保存结果表格
result_csv_path = os.path.join(output_dir, "categorical_analysis_results.csv")
result_df.to_csv(result_csv_path, index=False, encoding='utf-8-sig')
print(f"✅ 结果表格已保存到: {result_csv_path}")

# ========================
# 热力图可视化并保存
# ========================
plt.figure(figsize=(9, 5))
sns.heatmap(
    result_df.set_index('特征')[['统计值/相关系数']].T,
    annot=True, cmap='YlOrRd', cbar=True, fmt=".3f"
)
plt.title("类别特征与离职率(Attrition)的关系强度")
plt.xlabel("特征名称")

heatmap_path = os.path.join(output_dir, "attrition_relation_heatmap.png")
plt.tight_layout()
plt.savefig(heatmap_path, dpi=300)
plt.close()
print(f"✅ 热力图已保存到: {heatmap_path}")

# ========================
# 离职率柱状图（每个特征）
# ========================
for col in categorical_features:
    if data[col].nunique() == 1:
        continue
    plt.figure(figsize=(7, 4))
    sns.barplot(x=col, y=target, data=data, estimator=np.mean, palette="coolwarm")
    plt.title(f"{col} 不同类别的平均离职率")
    plt.xticks(rotation=30)
    plt.ylabel("平均离职率")
    plt.tight_layout()

    save_path = os.path.join(output_dir, f"{col}_attrition_bar.png")
    plt.savefig(save_path, dpi=300)
    plt.close()
    print(f"📊 已保存: {save_path}")

print("\n✅ 所有图片与结果文件已生成，保存在 ../fig 目录下。")